#ifndef GPU_FORMUNBALANCE_H
#define GPU_FORMUNBALANCE_H

#include <cuda_runtime.h>
#include <cusparse.h>
#include <cublas_v2.h>
#include <memory>
#include <vector>

// Forward declarations
class AnalysisModel;
class LinearSOE;
class FE_Element;
class DOF_Group;
class Vector;
class ID;

/**
 * GPU-accelerated formUnbalance implementation for OpenSees
 * Optimized for large-scale models (>10000 DOF)
 */
class GPU_FormUnbalance {
public:
    GPU_FormUnbalance(int maxElements, int maxDOFs, int maxNodes);
    ~GPU_FormUnbalance();

    // Main interface - GPU version of formUnbalance
    int formUnbalance(AnalysisModel* theAnalysisModel, LinearSOE* theSOE);
    
    // Performance monitoring
    double getLastComputeTime() const { return lastComputeTime; }
    double getLastTransferTime() const { return lastTransferTime; }
    void enableProfiling(bool enable) { profilingEnabled = enable; }

private:
    // GPU memory management
    struct GPUData {
        // Element data
        double* d_elementCoords;     // Element node coordinates
        int* d_elementConnectivity;  // Element-node connectivity
        int* d_elementTypes;         // Element type IDs
        double* d_elementProperties; // Material properties per element
        int* d_elementDOFMaps;       // Element DOF to global DOF mapping
        
        // Global arrays
        double* d_globalCoords;      // All node coordinates
        double* d_globalDisplacements; // Current displacements
        double* d_globalResidual;    // Global residual vector
        
        // Working arrays
        double* d_elementResiduals;  // Element residual vectors
        double* d_workSpace;         // Temporary workspace
        
        // Nodal data
        double* d_nodalLoads;        // Applied nodal loads
        int* d_nodalDOFMaps;         // Nodal DOF mappings
        
        // Sparse assembly data
        int* d_assemblyIndices;      // Assembly index mapping
        int* d_assemblyOffsets;      // Assembly offset array
    };

    // GPU streams for overlapping computation
    struct GPUStreams {
        cudaStream_t computeStream;
        cudaStream_t transferStream;
        cudaEvent_t computeEvent;
        cudaEvent_t transferEvent;
    };

    GPUData gpuData;
    GPUStreams streams;
    
    // CUDA libraries
    cublasHandle_t cublasHandle;
    cusparseHandle_t cusparseHandle;
    
    // Configuration
    int maxElements;
    int maxDOFs;
    int maxNodes;
    int currentElements;
    int currentDOFs;
    int currentNodes;
    
    // Performance monitoring
    bool profilingEnabled;
    double lastComputeTime;
    double lastTransferTime;
    cudaEvent_t startEvent, stopEvent;

    // Initialization and cleanup
    int initializeGPU();
    void cleanupGPU();
    
    // Data transfer methods
    int transferDataToGPU(AnalysisModel* theAnalysisModel);
    int transferResultFromGPU(LinearSOE* theSOE);
    
    // GPU computation kernels
    int computeElementResidualsGPU();
    int computeNodalUnbalanceGPU();
    int assembleGlobalResidualGPU();
    
    // Helper methods
    int updateElementData(AnalysisModel* theAnalysisModel);
    int updateNodalData(AnalysisModel* theAnalysisModel);
    bool checkGPUError(const char* operation);
    
    // Memory management helpers
    int allocateGPUMemory();
    void freeGPUMemory();
    int reallocateIfNeeded(int newElements, int newDOFs, int newNodes);
};

// CUDA kernel declarations
extern "C" {
    void launch_computeElementResiduals(
        const double* d_coords, const int* d_connectivity, const int* d_types,
        const double* d_properties, const double* d_displacements,
        const int* d_dofMaps, double* d_residuals,
        int numElements, int maxNodesPerElement, int maxDOFsPerElement,
        cudaStream_t stream
    );
    
    void launch_computeNodalUnbalance(
        const double* d_nodalLoads, const int* d_nodalDOFMaps,
        double* d_globalResidual, int numNodes, int maxDOFsPerNode,
        cudaStream_t stream
    );
    
    void launch_assembleGlobalResidual(
        const double* d_elementResiduals, const int* d_assemblyIndices,
        const int* d_assemblyOffsets, double* d_globalResidual,
        int numElements, int maxDOFsPerElement, cudaStream_t stream
    );
}

#endif // GPU_FORMUNBALANCE_H